"""
工具函数模块
"""

import cv2
import numpy as np
import logging
import json
import time
from typing import List, Dict, Any
from datetime import datetime

def setup_logging(log_level=logging.INFO):
    """设置日志系统"""
    logging.basicConfig(
        level=log_level,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler('traffic_system.log', encoding='utf-8'),
            logging.StreamHandler()
        ]
    )

def draw_visualization(frame, result):
    """绘制可视化结果"""
    if frame is None:
        return None
        
    vis_frame = frame.copy()
    
    # 绘制检测框和轨迹
    if 'tracks' in result:
        for track in result['tracks']:
            draw_track(vis_frame, track)
    
    # 绘制排队信息
    if 'queue_clusters' in result:
        draw_queue_clusters(vis_frame, result.get('queue_clusters', []))
    
    # 绘制统计信息
    draw_statistics(vis_frame, result)
    
    return vis_frame

def draw_track(frame, track):
    """绘制单个轨迹"""
    track_id = track.get('track_id', 0)
    bbox = track.get('bbox', [0, 0, 0, 0])
    trajectory = track.get('trajectory', [])
    
    # 绘制边界框
    x, y, w, h = bbox[:4]
    x1, y1 = int(x - w/2), int(y - h/2)
    x2, y2 = int(x + w/2), int(y + h/2)
    
    color = (0, 255, 0) if track.get('right_turn_intent', False) else (255, 0, 0)
    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
    
    # 绘制轨迹
    if len(trajectory) > 1:
        points = np.array(trajectory, dtype=np.int32)
        for i in range(1, len(points)):
            cv2.line(frame, tuple(points[i-1]), tuple(points[i]), color, 1)
    
    # 绘制ID和意图分数
    intent_score = track.get('intent_score', 0)
    text = f"ID:{track_id} S:{intent_score:.2f}"
    cv2.putText(frame, text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

def draw_queue_clusters(frame, clusters):
    """绘制排队集群"""
    colors = [
        (0, 255, 255),  # 黄色
        (255, 255, 0),  # 青色  
        (255, 0, 255),  # 粉色
        (0, 165, 255),  # 橙色
        (128, 0, 128)   # 紫色
    ]
    
    for i, cluster in enumerate(clusters):
        color = colors[i % len(colors)]
        
        # 绘制集群边界
        positions = []
        for vehicle in cluster:
            trajectory = vehicle.get('trajectory', [])
            if trajectory:
                positions.append(trajectory[-1])
        
        if positions:
            positions = np.array(positions, dtype=np.int32)
            
            # 计算凸包
            if len(positions) >= 3:
                hull = cv2.convexHull(positions)
                cv2.polylines(frame, [hull], True, color, 2)
            
            # 绘制集群中心
            center = np.mean(positions, axis=0).astype(int)
            cv2.circle(frame, tuple(center), 8, color, -1)
            cv2.putText(frame, f"Q{i}", (center[0]+10, center[1]), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

def draw_statistics(frame, result):
    """绘制统计信息"""
    # 背景矩形
    h, w = frame.shape[:2]
    cv2.rectangle(frame, (10, 10), (400, 150), (0, 0, 0), -1)
    cv2.rectangle(frame, (10, 10), (400, 150), (255, 255, 255), 1)
    
    # 统计信息
    stats = [
        f"帧号: {result.get('frame_id', 0)}",
        f"检测车辆: {result.get('detections', 0)}",
        f"跟踪目标: {result.get('tracks', 0)}", 
        f"右转车辆: {result.get('right_turn_vehicles', 0)}",
        f"排队集群: {result.get('queue_clusters', 0)}",
        f"当前长度: {result.get('current_length', 0):.1f}m",
        f"状态: {result.get('status', 'unknown')}"
    ]
    
    # 预测信息
    prediction = result.get('prediction', [])
    if prediction:
        stats.append(f"预测长度: {prediction[0]:.1f}m")
    
    # 绘制文本
    for i, text in enumerate(stats):
        y_pos = 30 + i * 20
        cv2.putText(frame, text, (20, y_pos), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

def calculate_performance_metrics(results):
    """计算性能指标"""
    if not results:
        return {}
    
    processing_times = []
    detection_counts = []
    queue_lengths = []
    
    for result in results:
        if result.get('status') == 'success':
            processing_times.append(result.get('processing_time', 0))
            detection_counts.append(result.get('detections', 0))
            
            current_length = result.get('current_length', 0)
            if current_length > 0:
                queue_lengths.append(current_length)
    
    metrics = {
        'total_frames': len(results),
        'successful_frames': len([r for r in results if r.get('status') == 'success']),
        'avg_processing_time': np.mean(processing_times) if processing_times else 0,
        'avg_detections': np.mean(detection_counts) if detection_counts else 0,
        'max_queue_length': np.max(queue_lengths) if queue_lengths else 0,
        'avg_queue_length': np.mean(queue_lengths) if queue_lengths else 0,
        'success_rate': len([r for r in results if r.get('status') == 'success']) / len(results)
    }
    
    return metrics

def save_results_to_json(results, filename):
    """保存结果到JSON文件"""
    try:
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(results, f, indent=2, ensure_ascii=False)
        logging.info(f"结果已保存到: {filename}")
    except Exception as e:
        logging.error(f"保存结果失败: {str(e)}")

def load_config_from_json(filename):
    """从JSON文件加载配置"""
    try:
        with open(filename, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        logging.error(f"加载配置失败: {str(e)}")
        return {}

def timestamp_to_datetime(timestamp):
    """时间戳转日期时间"""
    return datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S')

def create_video_writer(output_path, frame_size, fps=25):
    """创建视频写入器"""
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    return cv2.VideoWriter(output_path, fourcc, fps, frame_size)

def check_system_requirements():
    """检查系统要求"""
    requirements = {
        'opencv': cv2.__version__,
        'numpy': np.__version__,
        'torch': 'Not installed'
    }
    
    try:
        import torch
        requirements['torch'] = torch.__version__
    except ImportError:
        pass
    
    return requirements

def format_duration(seconds):
    """格式化时间间隔"""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    seconds = seconds % 60
    
    if hours > 0:
        return f"{hours}h {minutes}m {seconds:.1f}s"
    elif minutes > 0:
        return f"{minutes}m {seconds:.1f}s"
    else:
        return f"{seconds:.1f}s"

class PerformanceMonitor:
    """性能监控器"""
    
    def __init__(self, window_size=100):
        self.window_size = window_size
        self.processing_times = []
        self.memory_usage = []
        self.start_time = time.time()
    
    def record_processing_time(self, processing_time):
        """记录处理时间"""
        self.processing_times.append(processing_time)
        if len(self.processing_times) > self.window_size:
            self.processing_times.pop(0)
    
    def get_performance_stats(self):
        """获取性能统计"""
        if not self.processing_times:
            return {}
        
        current_time = time.time()
        elapsed_time = current_time - self.start_time
        
        stats = {
            'total_elapsed_time': elapsed_time,
            'total_frames_processed': len(self.processing_times),
            'current_fps': len(self.processing_times) / min(elapsed_time, 1),
            'avg_processing_time': np.mean(self.processing_times),
            'min_processing_time': np.min(self.processing_times),
            'max_processing_time': np.max(self.processing_times),
            'processing_time_std': np.std(self.processing_times)
        }
        
        return stats
    
    def print_performance_report(self):
        """打印性能报告"""
        stats = self.get_performance_stats()
        
        print("=" * 50)
        print("性能报告")
        print("=" * 50)
        for key, value in stats.items():
            if 'time' in key:
                print(f"{key}: {value:.4f}s")
            elif 'fps' in key:
                print(f"{key}: {value:.2f}")
            else:
                print(f"{key}: {value}")
        print("=" * 50)